#!/Pub/Users/wangyk/miniconda3/envs/pyscenic/bin/python3.10
import sys
import os
import argparse

wrp = '''
inferCNVpy命令行分析，利用infercnvpy进行CNV分析。在key_column列中，需要将上皮细胞分为多个亚群。比如说，一个seurat对象中包含上皮与成纤维细胞，上皮细胞分为0-9共计十个亚群，然后在本次分析中，会与对照细胞亚群进行inferCNVpy分析，识别0-9当中哪些上皮细胞亚群是恶性的。

命令参考：
inferCNV_py -s /Pub/Users/wangyk/project/Poroject/F240510002_OV_marco/out/SC_part/05.macro_sub_cellchat/ss_m_sub_anno_ready.rds -od Pub/Users/wangyk/project/Poroject/F240510002_OV_marco/tes2/ -k SubType -r TRM -w 12 -he 9
'''

parser = argparse.ArgumentParser(prog='inferCNV_py',
                                 description=wrp, formatter_class=argparse.RawTextHelpFormatter)

parser.add_argument('-h5', '--h5ad', help="h5ad文件路径。scanpy读取的单细胞分析h5ad对象",
                    default='', metavar='')
parser.add_argument('-s', '--seurat_obj', help="seurat对象文件路径。如果没有现成的h5ad对象，则尝试利用seurat对象的rds或rdata文件进行转化，会在结果路径中生成一个seurat_h5ad.h5ad文件",
                    default=False, metavar='')
parser.add_argument('-od', '--output_dir', help="结果输出路径",
                    default='./', metavar='')
parser.add_argument('-k', '--key_column',
                    help='用于识别细胞分群的列名，必须包含在AnnData的obs中，或者说seurat中meta.data中的某一列。比如用cell_type这一列的注释信息对细胞进行分群，其中包含0、1、2、3、T细胞亚群，然后使用T细胞亚群作为reference，对0、1、2、3亚群的细胞进行CNV识别。',
                    default='SubType', 
                    metavar='')
parser.add_argument('-r', '--reference', help="key_column列中，用于做对照的细胞类型。",
                    default='TRM', metavar='')
parser.add_argument('-w', '--width', help="数字，成图宽，默认为9",
                    default=9, metavar='', type=int)
parser.add_argument('-he', '--height', help="数字，成图高，默认为7",
                    default=7, metavar='', type=int)
parser.add_argument('-p', '--plot', help="逻辑值，对结果进行可视化，默认True",
                    default=True, metavar='')

args = parser.parse_args()



import infercnvpy as cnv
import pyarrow.feather as feather
import logging
import scanpy as sc

def check_file_existence(file_path: str):
    if not os.path.isfile(file_path):
        logger.error(f"未找到h5ad文件，请检查路径")
        sys.exit()

def ensure_directory_exists(directory_path: str):
    if not os.path.exists(directory_path):
        # logger.info(f"递归创建结果路径：{args.output_dir}")
        os.makedirs(directory_path)

ensure_directory_exists(args.output_dir)

formatter = logging.Formatter('%(asctime)s-%(levelname)s -%(message)s')
# 创建一个logger
logger = logging.getLogger(' infer_cnv_py ')
logger.setLevel(logging.INFO)

# 创建一个控制台处理器
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)

# 创建一个文件处理器
console_handler.setFormatter(formatter)
# 将处理器添加到logger
logger.addHandler(console_handler)


# trying convert seurat to h5ad
if bool(args.seurat_obj):
    r_file = "/Pub/Users/wangyk/Project_wangyk/Codelib_YK/some_scr/SC/inferCNV_py/seurat_2_h5ad.r"

    logger.info("转化Seurat对象到h5ad...")

    os.environ['R_HOME'] = "/Pub/Apps/Cellar/R/4.1.2/lib/R"
    os.system(f'{r_file} {args.seurat_obj} {args.output_dir}')

    args.h5ad = f'{args.output_dir}seurat_h5ad.h5ad'

# check files
check_file_existence(args.h5ad)

# loading h5ad obj
adata = sc.read_h5ad(args.h5ad)

if not args.key_column in adata.obs.columns.values:
    logger.error("指定列在AnnData.obs中不存在，请检查")
    sys.exit()


logger.info("Adding genomic position information to AnnData Obj ...")
cnv.io.genomic_position_from_gtf(
    "/Pub/Data/Data_Center/Annotation/gencode.v38.annotation.gtf", adata)

# cnv分析
cnv.tl.infercnv(
    adata,
    reference_key=args.key_column,
    reference_cat=args.reference,
    n_jobs=20
)
# 计算cnv score
cnv.tl.cnv_score(adata, groupby=args.key_column, key_added='cnv_score')

os.chdir(args.output_dir)
# save fig
cnv.pl.chromosome_heatmap(adata, groupby=args.key_column, figsize=(args.width, args.height), save='Fig_inferCNVpy.pdf')
# save score

adata.obs['cell'] = adata.obs.index
feather.write_feather(adata.obs, dest='Table_meta.data.feather')

if bool(args.plot):
    r_plot_file = "/Pub/Users/wangyk/Project_wangyk/Codelib_YK/some_scr/SC/inferCNV_py/inferCNV_py_plot_in_r.r"
    os.system(
        f"{r_plot_file} -k {args.key_column} -r {args.reference} -od {args.output_dir}")

logger.info("分析完成")


